############################################################################
# RE: several versions
#
############################################################################

rm(list=ls())
graphics.off()

library(data.table)
library(ggplot2)
library(haven)
library(xtable)
library(fixest)
library(mvtnorm)
library(Hmisc)

#############################################################################
# 1. Estimate different SE of beta coefficients
#############################################################################

# Bring data and FE results
load("Data_new/FE_results.RData")
rm(list=setdiff(ls(),list("dt","dt_res","dt_map")))
names(dt)
source("Code/0. CODE Auxiliary.R")
source("Code/6a. CODE RE.R")
time0=Sys.time()

f = feols(yield ~ precr + ddr:factor(fips) | factor(fips) + factor(year)^factor(state) ,data=dt)
s = summary(f,vcov="iid")
dt_slopes= slopes_dt(s,names_var="_iid")
setnames(dt_slopes,"coef_iid","coef")

###################################
# Add some summary variables
###################################

# Add corn-area as weights (sum per county)
dt_aux1 = dt[,lapply(.SD,sum),.SDcols="corn_area",by="fips"]
setnames(dt_aux1,c("fips","w"))
summary(dt_aux1)
# Mean of temperature (not weighted)
dt_aux2 = dt[,lapply(.SD,mean),.SDcols="ddr",by="fips"]
setnames(dt_aux2,c("fips","ddr_mean"))
summary(dt_aux2)
# Mean of temperature (weighted)
dt_aux3 = dt[,lapply(.SD,wtd.mean,weights=corn_area),.SDcols="ddr",by="fips"]
setnames(dt_aux3,c("fips","ddr_wmean"))
summary(dt_aux3)
# SD of temperature (not weighted)
dt_aux4 = dt[,lapply(.SD,sd),.SDcols="ddr",by="fips"]
setnames(dt_aux4,c("fips","ddr_sd"))
summary(dt_aux4)
# SD of temperature (weighted)
dt_aux5 = dt[,lapply(.SD,wtd.var,weights=corn_area),.SDcols="ddr",by="fips"]
dt_aux5[,ddr:=sqrt(ddr)]
setnames(dt_aux5,c("fips","ddr_wsd"))
summary(dt_aux5)

# Merge 
dim(dt_slopes)
dt_slopes = merge(dt_slopes,dt_aux1,by="fips",all=TRUE)
dt_slopes = merge(dt_slopes,dt_aux2,by="fips",all=TRUE)
dt_slopes = merge(dt_slopes,dt_aux3,by="fips",all=TRUE)
dt_slopes = merge(dt_slopes,dt_aux4,by="fips",all=TRUE)
dt_slopes = merge(dt_slopes,dt_aux5,by="fips",all=TRUE)
dim(dt_slopes)
summary(dt_slopes)

# Compare mean and sd of temperature with and without weights (over 95% correlation)
summary(dt_slopes[,.(ddr_mean,ddr_wmean,ddr_sd,ddr_wsd)])
cor(dt_slopes[,.(ddr_mean,ddr_wmean,ddr_sd,ddr_wsd)])

# Keep relevant objects
rm(list=setdiff(ls(),list("dt","dt_slopes","f","names_se","time0")))

############################################################################
# Prepare data for RE
############################################################################

data = copy(dt_slopes)

# Keep relevant variables and names
data=data[,.(fips,coef,se_iid,ddr_wmean,ddr_wsd,w)]
setnames(data,c("se_iid","ddr_wmean","ddr_wsd"),c("se","z","sd_z"))
summary(data)

# Data at the county level
dt_re = copy(data)
setnames(dt_re,"fips","id")
dim(dt_re)
summary(dt_re)

# Data at the county-year level with temperature and gamma to calculate marginal effects
dt_ext = dt[,.SD,.SDcols=c("fips","year","corn_area","ddr")]
setnames(dt_ext,c("id","t","w","ddr"))
dim(dt_ext)
summary(dt_ext)

# Source functions
source("Code/6a. CODE RE.R")

# 
dt_re[,w0:=w]
dt_ext[,w0:=w]

############################################################################
# RE (weighted)
############################################################################

dt_re[,w:=w0]
dt_ext[,w:=w0]
summary(dt_re$w)
summary(dt_ext$w)

# Constant prior
res_w1 = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="cte")

# Prior mean depends on temperature linearly, constant prior variance
res_w2 = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="mean_w")

# Prior mean depends on temperature linearly, prior variance is exponential of mean temperature
res_w3 = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="exp1")

# Prior mean depends on temperature linearly, prior variance is exponential of log mean temperature
# res_w4 = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="exp2")

# Prior mean depends on temperature linearly, prior variance is exponential of sd of temperature
# res_w5 = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="exp3")

# Prior mean depends on temperature linearly, prior variance is exponential of mean and sd of temperature
# res_w6 = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="exp4")

# Prior mean depends on temperature linearly, prior variance is binned
res_w7  = RE_func(data=dt_re,data_ext=dt_ext,min0=10,max0=200,model="bins")

time1=Sys.time()
save.image("Data_new/RE_results.RData")

# Check prior variance of beta are positive
summary(res_w1$dt$var0)
summary(res_w2$dt$var0)
summary(res_w3$dt$var0)
#summary(res_w4$dt$var0)
#summary(res_w5$dt$var0)
#summary(res_w6$dt$var0)
summary(res_w7$dt$var0)

############################################################################
# Tabulate results
############################################################################

rm(list=ls())
load("Data_new/RE_results.RData")
save_results = "Results"
source("Code/0. CODE Auxiliary.R")

#######################################
# Posterior means
#######################################
### Distribution of posterior means

### Uncorrelated RE
PM1 = copy(res_w1$dt_ext)
setnames(PM1,c("id","t"),c("fips","year"))

### Correlated RE only mean
PM2 = copy(res_w2$dt_ext)
setnames(PM2,c("id","t"),c("fips","year"))

### Correlated RE where the variance also depends on temperature (exp(x))
PM3 = copy(res_w3$dt_ext)
setnames(PM3,c("id","t"),c("fips","year"))

### Correlated RE where the variance also depends on temperature (3 bins)
PM4 = copy(res_w7$dt_ext)
setnames(PM4,c("id","t"),c("fips","year"))

#######################################
# Keep descriptive tables
#######################################

### Uncorrelated RE
ss  = res_w1$table_mg

### Correlated RE only mean
ss  = cbind(ss,res_w2$table_mg)

### Correlated RE where the variance also depends on temperature (exp(x))
ss  = cbind(ss,res_w3$table_mg)

### Correlated RE where the variance also depends on temperature (3 bins)
ss  = cbind(ss,res_w7$table_mg)

tab_RE = ss
tab_RE

############################################
# Table RE
############################################

# Column of posterior densities
s = tab_RE[1:(dim(tab_RE)[1]-1),]
s

# Column of distribution of posterior mean when only prior mean depends on x
s2 = sum_table(data=PM2,name_vars="mu1",save_name="",weight_var = "w")
s = cbind(s,s2)
s

l=dim(s)[2]
dig=rep(3,l+1)
print(xtable(s,type="latex",digits=dig,display=c("s",rep("f",l))),hline.after = NULL,
      file=paste0(paste0(save_results,"/Summary_table_RE"),".tex"),include.rownames = TRUE,include.colnames=FALSE,
      sanitize.text.function = function(x){x},only.contents = TRUE)
